import pandas as pd
import matplotlib.pyplot as plt


def co2_gdp_plot():
    # 读取数据集
    df_climate = pd.read_csv('ClimateChange.csv')

    # 筛选需要的列
    df_co2 = df_climate[df_climate['Series code'] == 'EN.ATM.CO2E.KT']
    df_gdp = df_climate[df_climate['Series code'] == 'NY.GDP.MKTP.CD']

    # 填充缺失值
    df_co2 = df_co2.fillna(method='ffill', axis=1).fillna(method='bfill', axis=1)
    df_gdp = df_gdp.fillna(method='ffill', axis=1).fillna(method='bfill', axis=1)

    # 按国家分组求和
    df_co2_sum = df_co2.iloc[:, 5:].sum(axis=1)
    df_gdp_sum = df_gdp.iloc[:, 5:].sum(axis=1)

    # 归一化处理
    df_co2_norm = (df_co2_sum - df_co2_sum.min()) / (df_co2_sum.max() - df_co2_sum.min())
    df_gdp_norm = (df_gdp_sum - df_gdp_sum.min()) / (df_gdp_sum.max() - df_gdp_sum.min())

    # 创建子图对象
    fig, ax = plt.subplots()

    # 绘制曲线图
    ax.plot(df_co2_norm, label='CO2-SUM')
    ax.plot(df_gdp_norm, label='GDP-SUM')

    # 设置标题和标签
    ax.set_title('GDP-CO2')
    ax.set_xlabel('Countries')
    ax.set_ylabel('Values')

    # 设置横坐标刻度
    ax.set_xticks([0, 1, 2, 3, 4])
    ax.set_xticklabels(['CHN', 'USA', 'RUS', 'FRA', 'GBR'])

    # 添加图例
    ax.legend()

    # 显示图形
    plt.show()

    # 返回中国归一化后的 CO2 和 GDP 数据
    china_co2_norm = round(df_co2_norm[df_co2['Country code'] == 'CHN'].values[0], 3)
    china_gdp_norm = round(df_gdp_norm[df_gdp['Country code'] == 'CHN'].values[0], 3)

    return [china_co2_norm, china_gdp_norm]
